package _07_映射;

import _01_第一季._00_公共接口.Map;

import java.util.Comparator;
import java.util.LinkedList;
import java.util.Queue;

/*
    底层由红黑树实现的TreeMap映射,判断元素K支持可比较性
 */
public class TreeMap<K,V> implements Map<K,V>{
    //红黑树颜色节点的定义 红为假  黑为真
    private static final boolean RED = false;
    private static final boolean BLACK = true;
    private int size;
    private Node<K,V> root; //根节点
    private Comparator<K> comparator;

    public TreeMap(){
        this(null);
    }

    public TreeMap(Comparator<K> comparator){
        this.comparator = comparator;
    }
    
    private static class Node<K,V>{
        K key;
        V value;
        boolean color = RED;
        Node<K,V> left;
        Node<K,V> right;
        Node<K,V> parent;
        public Node(){
            this(null,null,null);
        }
        public Node(K key,V value,Node<K,V> parent){
            this.key = key;
            this.value = value;
            this.parent = parent;
        }

        public boolean isLeaf(){
            return left == null && right == null;
        }
        public boolean hasTwoKids(){
            return left!=null && right != null;
        }
        public boolean isLeftChild(){
            return parent!=null && this == parent.left;
        }
        public boolean isRightChild(){
            return parent!=null && this == parent.right;
        }
        //返回兄弟节点
        public Node<K,V> sibling(){
            if(isLeftChild()){
                return parent.right;
            }
            if(isRightChild()){
                return parent.left;
            }
            return null;
        }
    }

    public int size(){
        return size;
    }

    public boolean isEmpty(){
        return size == 0;
    }

    public void clear(){
        root = null;
        size = 0;
    }

    /**
     * 添加的方法
     * @param key : 键
     * @param value : 值
     * @return : 如果有重复的值，就返回覆盖之前的值
     */
    @Override
    public V put(K key, V value) {
        keyNotNullCheck(key);
        //添加的第一个节点 根节点
        if(root == null){
            root = new Node<>(key,value,null);
            size++;
            //添加之后进行调整
            afterPut(root);
            return null;
        }

        //添加的不是第一个节点
        Node<K,V> node = root; //从根节点开始比较
        Node<K,V> parent = root; //记录父节点
        int cmp = 0; //记录比较结果
        do{
            cmp = compare(key,node.key); //每一轮进行比较
            parent = node; //记录父节点
            if(cmp > 0){
                node = node.right;
            }else if(cmp < 0){
                node = node.left;
            }else{ //相等的情况
                V ret = node.value;
                node.key = key; //直接覆盖即可
                node.value = value;
                return ret;
            }
        }while (node!=null);

        Node<K,V> newNode = new Node<>(key,value,parent);
        //根据方向进行插入节点
        if(cmp > 0){
            parent.right = newNode;
        }else{
            parent.left = newNode;
        }
        size++;

        //添加之后进行平衡调整
        afterPut(newNode);
        return null;
    }

    /**
     * 添加之后的修复
     * @param node
     */
    private void afterPut(Node<K, V> node) {
        Node<K,V> parent = node.parent;
        //添加的是根节点或者上溢到达了根节点，直接染成黑色返回即可
        if(parent == null){
            black(node);
            return;
        }

        //情况1 父节点是黑色
        if(isBlack(parent)) return;

        //拿到叔父节点，根据判定条件
        Node<K,V> uncle = parent.sibling();
        //祖父节点
        Node<K,V> grand = parent.parent;

        //情况3 叔父节点是红色 上溢
        if(isRed(uncle)){
            black(parent);
            black(uncle);
            afterPut(red(grand)); //递归处理
            return;
        }

        //情况2 叔父节点不是红色 旋转
        if(parent.isLeftChild()){
            red(grand);
            if(node.isLeftChild()){ //LL
                black(parent);
            }else{ //LR
                black(node);
                rotateLeft(parent);
            }
            rotateRight(grand);
        }else{
            red(grand);
            if(node.isLeftChild()){ //RL
                black(node);
                rotateRight(parent);
            }else{ //RR
                black(parent);
            }
            rotateLeft(grand);
        }
    }

    //根据键找到对应的值并且返回，如果为空就不返回
    @Override
    public V get(K key) {
        Node<K, V> node = node(key);
        return node != null ? node.value : null;
    }

    //根据元素找到元素所在节点
    private Node<K,V> node(K key){
        keyNotNullCheck(key);
        Node<K,V> node = root;
        while (node != null){
            int cmp = compare(key,node.key);
            if(cmp == 0) return node;
            if(cmp > 0){
                node = node.right;
            }else{
                node = node.left;
            }
        }
        return null;
    }

    //返回待删除节点键对应的值
    @Override
    public V remove(K key) {
        return remove(node(key));
    }

    private V remove(Node<K,V> node){
        if(node == null) return null;
        V oldValue = node.value;
        size--;

        //首先考虑删除度为2的节点，找到其前驱节点或者后继节点，将其转换为删除度为1或者度为0的节点
        if(node.hasTwoKids()){
            Node<K,V> pre = predecessor(node);
            node.key = pre.key;
            node.value = pre.value;
            node = pre; //下面统一进行处理
        }

        //删除度为1或者度为0的节点
        Node<K,V> replacement = node.left != null ? node.left : node.right;
        if(replacement != null){ //删除度为1的节点
            replacement.parent = node.parent;
            if(node.parent == null){ //删除度为1的根节点
                root = replacement;
            } else if(node.isLeftChild()){
                node.parent.left = replacement;
            }else{
                node.parent.right = replacement;
            }
            //删除度为1的节点以后的调整
            afterRemove(replacement);
        }else if(node.parent == null){ //删除度为0的根节点
            root = null;
            //删除度为0的根节点以后的调整
            afterRemove(node);
        }else{ //删除度为0的叶子节点
            if(node.isLeftChild()){
                node.parent.left = null;
            }else{
                node.parent.right = null;
            }
            //删除度为0的叶子节点以后的调整
            afterRemove(node);
        }
        return oldValue;
    }

    /**
     * 删除之后的修复
     * @param node
     */
    private void afterRemove(Node<K, V> node) {
        if(isRed(node)){
            black(node);
            return;
        }

        //删除的是根节点
        Node<K,V> parent = node.parent;
        if(parent == null) return;

        //删除黑色叶子节点
        //需要找到兄弟节点，但是兄弟节点在之前的删除逻辑中父节点的那根线已经断了，无法通过node.sibling()找到兄弟节点
        //逆向思维 判断被删除的节点在左还是右
        boolean left = parent.left == null || node.isLeftChild(); //下溢后产生的下溢现象，可以用第二个条件判断
        Node<K,V> sibling = left ? parent.right : parent.left;

        if(left){ //被删除的节点在左边 兄弟节点在右边
            //删除的节点在左边和删除的节点在右边逻辑完全一样，只是方向相反
            if(isRed(sibling)){
                black(sibling);
                red(parent);
                rotateLeft(parent);
                sibling = parent.right;
            }
            if(isBlack(sibling.left) && isBlack(sibling.right)){
                //没有，借不了
                boolean parentBlack = isBlack(parent);
                black(parent);
                red(sibling);
                if(parentBlack){ //父节点是黑色的情况，导致父节点也会下溢
                    afterRemove(parent);
                }
            }else{
                //有至少一个，可以借
                if(isBlack(sibling.right)){
                    rotateRight(sibling);
                    sibling = parent.right; //旋转以后 兄弟节点也变了
                }
                color(sibling,colorOf(parent));
                black(sibling.right);
                black(parent);
                rotateLeft(parent);
            }

        }else{ //被删除的在右边 兄弟节点在左边

            //先考虑兄弟节点是红色的情况，将其转成兄弟节点为黑色的情况,下面就能统一处理
            if(isRed(sibling)){
                black(sibling);
                red(parent);
                rotateRight(parent);
                sibling = parent.left;
            }

            //其次考虑兄弟节点是黑色的情况
            //看兄弟节点是否有一个红色的子节点可以借
            if(isBlack(sibling.left) && isBlack(sibling.right)){
                //没有，借不了
                boolean parentBlack = isBlack(parent);
                black(parent);
                red(sibling);
                if(parentBlack){ //父节点是黑色的情况，导致父节点也会下溢
                    afterRemove(parent);
                }
            }else{
                //有至少一个，可以借
                //先处理需要双旋的情况，兄弟节点的左节点是黑色，将兄弟节点左旋转，后面即可统一处理
                if(isBlack(sibling.left)){
                    rotateLeft(sibling);
                    sibling = parent.left; //旋转以后 兄弟节点也变了
                }
                //统一对父节点进行右旋转,再将中心节点染成和父节点一样的颜色，左右节点染成黑色
                //中心节点其实也就是兄弟节点
                color(sibling,colorOf(parent));
                black(sibling.left);
                black(parent);
                rotateRight(parent);
            }
        }
    }

    //是否包含某个key
    @Override
    public boolean containsKey(K key) {
        return node(key) != null;
    }

    //是否包含某个值，使用层序遍历
    @Override
    public boolean containsValue(V value) {
        if(root == null) return false;
        Queue<Node<K,V>> queue = new LinkedList<>();
        queue.offer(root);
        while(!queue.isEmpty()){
            Node<K, V> node = queue.poll();
            if(valEquals(value,node.value)) return true;

            if(node.left!=null){
                queue.offer(node.left);
            }
            if(node.right!=null){
                queue.offer(node.right);
            }
        }
        return false;
    }

    //内部比较方法
    private boolean valEquals(V v1,V v2){
        return (v1 == v2) || (v1 != null && v1.equals(v2));
    }

    //遍历的方法，使用中序遍历
    @Override
    public void traversal(Visitor<K, V> visitor) {
        if(visitor == null) return;
        traversal(root,visitor);
    }

    //内部中序遍历的方法
    private void traversal(Node<K,V> node,Visitor<K, V> visitor) {
        if(node == null || visitor.stop) return;

        traversal(node.left,visitor);
        if(visitor.stop) return; //停止遍历
        visitor.visit(node.key,node.value);
        traversal(node.right,visitor);
    }

    //旋转之后更新父节点以及高度的统一操作
    private void afterRotate(Node<K,V> grand, Node<K,V> parent, Node<K,V> child){
        //更新p的父节点
        parent.parent = grand.parent;
        if(grand.isLeftChild()){
            grand.parent.left = parent;
        }else if(grand.isRightChild()){
            grand.parent.right = parent;
        }else{
            root = parent;
        }

        //更新child的父节点
        if(child!=null){
            child.parent = grand;
        }

        //更新grand的父节点
        grand.parent = parent;

    }

    //左旋转
    private void rotateLeft(Node<K,V> grand){
        Node<K,V> parent = grand.right;
        Node<K,V> child = parent.left;

        //换线
        grand.right = child;
        parent.left = grand;

        //旋转之后更新父节点以及高度的操作
        afterRotate(grand,parent,child);

    }

    //右旋转
    private void rotateRight(Node<K,V> grand){
        Node<K,V> parent = grand.left;
        Node<K,V> child = parent.right;
        grand.left = child;
        parent.right = grand;

        //旋转之后更新父节点以及高度的操作
        afterRotate(grand,parent,child);

    }


    //找到该节点的前驱节点:中序遍历时的前一个节点
    private Node<K,V> predecessor(Node<K,V> node){
        if(node == null) return null;
        Node<K,V> pre = node.left;
        //情况1
        if(pre!=null){
            while(pre.right!=null){
                pre = pre.right;
            }
            return pre;
        }

        //情况2 终止条件 前驱节点在父节点的右子树中
        while(node.parent!=null && node == node.parent.left){
            node = node.parent;
        }
        return node.parent;
    }

    //找到该节点的后继节点:中序遍历时的后一个节点
    private Node<K,V> successor(Node<K,V> node){
        Node<K,V> suc = node.right;

        //情况1
        if(suc!=null){
            while(suc.left!=null){
                suc = suc.left;
            }
            return suc;
        }

        //情况2 终止条件 前驱节点在父节点的左子树中
        while(node.parent!=null && node == node.parent.right){
            node = node.parent;
        }
        return node.parent;
    }

    //内部比较方法
    private int compare(K k1,K k2){
        //如果该BST实现了comparator接口，那么内部比较方法就是用比较器
        if(comparator != null){
            return comparator.compare(k1, k2);
        }
        //否则就强制转换成自然排序接口进行比较
        return ((Comparable<K>)k1).compareTo(k2);
    }

    //检测Key内部不能为null
    private void keyNotNullCheck(K key){
        if(key == null){
            throw new IllegalArgumentException("key不能为null");
        }
    }

    /**
     * 对节点染色的方法
     * @param node 待染色的节点
     * @param color 颜色
     * @return 返回染色后的节点
     */
    private Node<K,V> color(Node<K,V> node, boolean color){
        if(node == null) return null;
        node.color = color;
        return node;
    }

    /**
     * 将节点染成红色
     * @param node 带染色的节点
     * @return 返回染成红色的节点
     */
    private Node<K,V> red(Node<K,V> node){
        return color(node,RED);
    }

    /**
     * 将节点染成黑色
     * @param node 带染色的节点
     * @return 返回染成黑色的节点
     */
    private Node<K,V> black(Node<K,V> node){
        return color(node,BLACK);
    }

    /**
     * 判断节点颜色的方法
     * @param node
     * @return 返回节点的颜色
     */
    private boolean colorOf(Node<K,V> node){
        return node == null ? BLACK : node.color;
    }

    //判断节点颜色是否是红色
    private boolean isRed(Node<K,V> node){
        return colorOf(node) == RED;
    }

    //判断节点颜色是否是黑色
    private boolean isBlack(Node<K,V> node){
        return colorOf(node) == BLACK;
    }
}
